
import numpy as np
from scipy.linalg import expm
import matplotlib.pyplot as plt
from pathlib import Path  # 新增

# 保存目录（新增）
SAVE_DIR = Path(r"F:\python_Summer tasks\figures")
SAVE_DIR.mkdir(parents=True, exist_ok=True)

# =================== 全局配置 ===================
CONFIG = {
    # --- 系统量子比特数 ---
    "num_qubits": 5,
    # --- 时间网格 ---
    "t_start": 0.0,
    "t_end": 3.0,
    "num_time_samples":1000,

    # --- 子空间规模  p  ---
    "num_dict": 5,
    "dict_thetas_mode": "linspace",

    # --- 损失权重 ---
    "w_data": 8,
    "w_phys": 0.001,
    "w_norm": 0.001,
    "lambda_reg": 1e-5,

    # --- 优化超参 ---
    "lr_eq12": 5e-4,
    "lr_eq11": 2e-4,
    "num_iters": 400,

    # --- 时间导数离散 ---
    "diff_scheme": "central",

    # --- 哈密顿量参数（复杂度开关都在这） ---
    # basic / nn / full / random
    "ham_mode": "random",
    "H_local_Z": 0.0005,
    "H_local_X": 0.56,
    "two_body_scale": 0.005,
    "time_mod_amp": 0.06,
    "ham_random_seed": 41,
    "use_global_X": True,
    "H_global_X_amp0": 0.21,
    "H_global_X_amp1": 0.32,

    # -------- 噪声实验 --------
    "noise_sigma_list": [0,0],
    "num_trials_per_sigma": 8,
    "base_seed": 1234,
    # 噪声注入方式: "frozen" = 优化前抽一次；"per_iter" = 每一步都抽
    "noise_injection_mode": "per_iter",

    # -------- 想看收敛/可观测量的那一档噪声 --------
    # 这档会额外记录 loss 和 observable
    "plot_sigma": 0.0,
}
# ==================================================


# ============== 工具 ==============
def kron_n(*mats):
    out = np.array([[1]], dtype=complex)
    for M in mats:
        out = np.kron(out, M)
    return out


def pauli_mats():
    X = np.array([[0, 1],
                  [1, 0]], dtype=complex)
    Y = np.array([[0, -1j],
                  [1j, 0]], dtype=complex)
    Z = np.array([[1, 0],
                  [0, -1]], dtype=complex)
    I = np.eye(2, dtype=complex)
    return X, Y, Z, I


# ============== 1. H(t)（复杂版） ==============
def H_of_t(t, cfg):
    n = cfg["num_qubits"]
    X1, Y1, Z1, I1 = pauli_mats()
    dim = 2 ** n
    H = np.zeros((dim, dim), dtype=complex)

    # 单体
    hZ = cfg.get("H_local_Z", 0.0)
    hX = cfg.get("H_local_X", 0.0)
    for q in range(n):
        if hZ != 0.0:
            ops = [Z1 if k == q else I1 for k in range(n)]
            H += hZ * kron_n(*ops)
        if hX != 0.0:
            ops = [X1 if k == q else I1 for k in range(n)]
            H += hX * kron_n(*ops)

    ham_mode = cfg.get("ham_mode", "basic")
    two_body_scale = cfg.get("two_body_scale", 0.3)
    time_mod_amp = cfg.get("time_mod_amp", 0.0)

    rng = None
    if ham_mode == "random":
        rng = np.random.default_rng(cfg.get("ham_random_seed", 42))

    # 两体
    if ham_mode == "nn":
        for i in range(n - 1):
            # ZZ
            ops = [(Z1 if k == i or k == i+1 else I1) for k in range(n)]
            H += (two_body_scale + time_mod_amp * np.cos(t)) * kron_n(*ops)
            # XX
            ops = [(X1 if k == i or k == i+1 else I1) for k in range(n)]
            H += 0.5 * two_body_scale * kron_n(*ops)

    elif ham_mode == "full":
        for i in range(n):
            for j in range(i+1, n):
                ops = [(Z1 if k == i or k == j else I1) for k in range(n)]
                H += (two_body_scale + time_mod_amp * np.cos(t)) * kron_n(*ops)

    elif ham_mode == "random":
        for i in range(n):
            for j in range(i+1, n):
                a_zz = two_body_scale * rng.normal()
                a_xx = two_body_scale * rng.normal()
                a_xy = two_body_scale * 0.5 * rng.normal()

                # ZZ
                ops = [(Z1 if k == i or k == j else I1) for k in range(n)]
                H += (a_zz + time_mod_amp * np.cos(t)) * kron_n(*ops)

                # XX
                ops = [(X1 if k == i or k == j else I1) for k in range(n)]
                H += a_xx * kron_n(*ops)

                # (XY + YX)
                ops_xy, ops_yx = [], []
                for k in range(n):
                    if k == i:
                        ops_xy.append(X1); ops_yx.append(Y1)
                    elif k == j:
                        ops_xy.append(Y1); ops_yx.append(X1)
                    else:
                        ops_xy.append(I1); ops_yx.append(I1)
                H += a_xy * (kron_n(*ops_xy) + kron_n(*ops_yx))

    # 全局 X...X
    if cfg.get("use_global_X", True):
        X_ops = [X1 for _ in range(n)]
        X_all = kron_n(*X_ops)
        a0 = cfg.get("H_global_X_amp0", 0.6)
        a1 = cfg.get("H_global_X_amp1", 0.3)
        H += (a0 + a1 * np.cos(t)) * X_all

    return H


def exact_time_evolution(t_grid, psi0, cfg):
    dt = t_grid[1] - t_grid[0]
    states = [psi0]
    for k in range(len(t_grid) - 1):
        t_mid = 0.5 * (t_grid[k] + t_grid[k+1])
        U = expm(-1j * H_of_t(t_mid, cfg) * dt)
        states.append(U @ states[-1])
    return np.stack(states, axis=0)


# ============== 2. n-qubit PQC ==============
def cnot_on_n_qubits(state, n, control, target):
    dim = 2**n
    new_state = state.copy()
    for basis in range(dim):
        if (basis >> (n-1-control)) & 1:
            flipped = basis ^ (1 << (n-1-target))
            new_state[flipped] = state[basis]
        else:
            new_state[basis] = state[basis]
    return new_state


def feature_state_nqubit(t, theta, cfg):
    n = cfg["num_qubits"]
    dim = 2**n
    state = np.zeros(dim, dtype=complex); state[0] = 1.0
    X1, Y1, Z1, I1 = pauli_mats()
    for q in range(n):
        Rz_theta = np.array([[np.exp(-1j*theta/2), 0],
                             [0, np.exp(1j*theta/2)]], dtype=complex)
        Rx_t = np.array([[np.cos(t/2), -1j*np.sin(t/2)],
                         [-1j*np.sin(t/2), np.cos(t/2)]], dtype=complex)
        Rz_t2 = np.array([[np.exp(-1j*t/4), 0],
                          [0, np.exp(1j*t/4)]], dtype=complex)
        for gate in (Rz_theta, Rx_t, Rz_t2):
            ops = [gate if k == q else I1 for k in range(n)]
            Uq = kron_n(*ops)
            state = Uq @ state
    for c in range(n-1):
        state = cnot_on_n_qubits(state, n, c, c+1)
    return state


def build_dict_states(t_grid, cfg):
    ndict = cfg["num_dict"]
    if cfg["dict_thetas_mode"] == "linspace":
        thetas = np.linspace(0, 2*np.pi, ndict, endpoint=False)
    else:
        thetas = 2*np.pi * np.random.rand(ndict)
    ntime = len(t_grid)
    dim = 2**cfg["num_qubits"]
    out = np.zeros((ndict, ntime, dim), dtype=complex)
    for i, th in enumerate(thetas):
        for j, t in enumerate(t_grid):
            out[i, j] = feature_state_nqubit(t, th, cfg)
    return out


# ============== 3. build_core (clean) ==============
def build_core(t_grid, dict_states, psi0, cfg):
    ndict, ntime, dim = dict_states.shape
    dt = t_grid[1] - t_grid[0]

    D0 = dict_states[:, 0, :].T  # (dim, ndict)

    phys_list = []
    for j in range(ntime):
        t = t_grid[j]
        if cfg["diff_scheme"] == "central":
            if j == 0:
                dpsi = (dict_states[:, 1, :] - dict_states[:, 0, :]) / dt
            elif j == ntime - 1:
                dpsi = (dict_states[:, -1, :] - dict_states[:, -2, :]) / dt
            else:
                dpsi = (dict_states[:, j+1, :] - dict_states[:, j-1, :]) / (2*dt)
        elif cfg["diff_scheme"] == "forward":
            if j == ntime - 1:
                dpsi = (dict_states[:, j, :] - dict_states[:, j-1, :]) / dt
            else:
                dpsi = (dict_states[:, j+1, :] - dict_states[:, j, :]) / dt
        else:  # backward
            if j == 0:
                dpsi = (dict_states[:, 1, :] - dict_states[:, 0, :]) / dt
            else:
                dpsi = (dict_states[:, j, :] - dict_states[:, j-1, :]) / dt

        Hj = H_of_t(t, cfg)
        psi_j = dict_states[:, j, :]
        Hpsi = (Hj @ psi_j.T).T
        Mj = 1j * dpsi - Hpsi
        phys_list.append(Mj.T)

    K_list = []
    for j in range(ntime):
        psi_j = dict_states[:, j, :]
        K = psi_j @ np.conj(psi_j.T)
        K_list.append(K)

    return D0, phys_list, K_list


# ============== 4. 噪声注入 ==============
def add_noise_to_core(D0, phys_list, K_list, sigma):
    if sigma == 0.0:
        return D0, phys_list, K_list

    def noisy(x):
        noise = sigma * (np.random.randn(*x.shape) + 1j*np.random.randn(*x.shape))
        return x + np.abs(x) * noise

    D0_n = noisy(D0)
    phys_n = [noisy(Mj) for Mj in phys_list]

    K_n = []
    for K in K_list:
        E = noisy(K)
        E = 0.5 * (E + E.conj().T)
        K_n.append(E)
    return D0_n, phys_n, K_n


# ============== 5. 两个loss ==============
def loss_and_grad_eq12(alpha, D0, phys_list, K_list, psi0, cfg):
    ntime = len(phys_list)
    lam = cfg["lambda_reg"]

    r_data = D0 @ alpha - psi0
    L_data = cfg["w_data"] * np.vdot(r_data, r_data).real
    g_data = 2 * cfg["w_data"] * (D0.conj().T @ r_data)

    L_phys = 0.0
    g_phys = np.zeros_like(alpha)
    for Mj in phys_list:
        rj = Mj @ alpha
        L_phys += np.vdot(rj, rj).real
        g_phys += 2 * (Mj.conj().T @ rj)
    L_phys = cfg["w_phys"] * (L_phys / ntime)
    g_phys = cfg["w_phys"] * (g_phys / ntime)

    reg_val = 0.0
    reg_grad = np.zeros_like(alpha)
    for K in K_list:
        s = np.vdot(alpha, K @ alpha)
        reg_val += s.real
        reg_grad += 2 * (K.real @ alpha)
    reg_val = lam * reg_val / ntime
    reg_grad = lam * reg_grad / ntime

    L_total = L_data + L_phys + reg_val
    g_total = g_data + g_phys + reg_grad
    return L_total, g_total


def loss_and_grad_eq11(alpha, D0, phys_list, K_list, psi0, cfg):
    ntime = len(phys_list)
    lam = cfg["lambda_reg"]

    r_data = D0 @ alpha - psi0
    L_data = cfg["w_data"] * np.vdot(r_data, r_data).real
    g_data = 2 * cfg["w_data"] * (D0.conj().T @ r_data)

    L_phys = 0.0
    g_phys = np.zeros_like(alpha)
    for Mj in phys_list:
        rj = Mj @ alpha
        L_phys += np.vdot(rj, rj).real
        g_phys += 2 * (Mj.conj().T @ rj)
    L_phys = cfg["w_phys"] * (L_phys / ntime)
    g_phys = cfg["w_phys"] * (g_phys / ntime)

    L_norm = 0.0
    g_norm = np.zeros_like(alpha)
    for K in K_list:
        s = np.vdot(alpha, K @ alpha).real
        diff = s - 1.0
        L_norm += diff * diff
        g_norm += 4 * diff * (K.real @ alpha)
    L_norm = cfg["w_norm"] * (L_norm / ntime)
    g_norm = cfg["w_norm"] * (g_norm / ntime)

    reg_val = 0.0
    reg_grad = np.zeros_like(alpha)
    for K in K_list:
        s = np.vdot(alpha, K @ alpha)
        reg_val += s.real
        reg_grad += 2 * (K.real @ alpha)
    reg_val = lam * reg_val / ntime
    reg_grad = lam * reg_grad / ntime

    L_total = L_data + L_phys + L_norm + reg_val
    g_total = g_data + g_phys + g_norm + reg_grad
    return L_total, g_total


# ============== 6. 重建 & 可观测量 ==============
def reconstruct(dict_states, alpha):
    nd, nt, dim = dict_states.shape
    out = np.zeros((nt, dim), dtype=complex)
    for j in range(nt):
        out[j] = np.tensordot(alpha, dict_states[:, j, :], axes=(0, 0))
    return out


def build_observable_Zsum(cfg):
    """Σ Z_i"""
    n = cfg["num_qubits"]
    X1, Y1, Z1, I1 = pauli_mats()
    dim = 2**n
    Z_sum = np.zeros((dim, dim), dtype=complex)
    for q in range(n):
        ops = [Z1 if k == q else I1 for k in range(n)]
        Z_sum += kron_n(*ops)
    return Z_sum


def exp_O(states, O):
    return np.einsum('ti,ij,tj->t', np.conj(states), O, states)


# ============== 7. 单次实验 ==============
def run_once(sigma, dict_states, D0_clean, phys_clean, K_clean,
             psi0, exact_states, cfg,
             record_history=False):
    ndict = cfg["num_dict"]

    hist12 = []
    hist11 = []

    # -------- 式(12) --------
    alpha12 = np.zeros(ndict, dtype=complex)
    for it in range(cfg["num_iters"]):
        if cfg["noise_injection_mode"] == "per_iter":
            D0_n, phys_n, K_n = add_noise_to_core(D0_clean, phys_clean, K_clean, sigma)
        else:
            if it == 0:
                D0_n, phys_n, K_n = add_noise_to_core(D0_clean, phys_clean, K_clean, sigma)
        L, g = loss_and_grad_eq12(alpha12, D0_n, phys_n, K_n, psi0, cfg)
        if record_history:
            hist12.append(L)
        alpha12 -= cfg["lr_eq12"] * g
    approx12 = reconstruct(dict_states, alpha12)
    err12 = np.linalg.norm(approx12 - exact_states, axis=1).mean()

    # -------- 式(11) --------
    alpha11 = np.zeros(ndict, dtype=complex)
    for it in range(cfg["num_iters"]):
        if cfg["noise_injection_mode"] == "per_iter":
            D0_n, phys_n, K_n = add_noise_to_core(D0_clean, phys_clean, K_clean, sigma)
        else:
            if it == 0:
                D0_n, phys_n, K_n = add_noise_to_core(D0_clean, phys_clean, K_clean, sigma)
        L, g = loss_and_grad_eq11(alpha11, D0_n, phys_n, K_n, psi0, cfg)
        if record_history:
            hist11.append(L)
        alpha11 -= cfg["lr_eq11"] * g
    approx11 = reconstruct(dict_states, alpha11)
    err11 = np.linalg.norm(approx11 - exact_states, axis=1).mean()

    if record_history:
        return err12, err11, approx12, approx11, hist12, hist11
    else:
        return err12, err11


# ============== 8. main ==============
if __name__ == "__main__":
    cfg = CONFIG
    n = cfg["num_qubits"]
    dim = 2**n

    t_grid = np.linspace(cfg["t_start"], cfg["t_end"], cfg["num_time_samples"])
    psi0 = np.zeros(dim, dtype=complex); psi0[0] = 1.0

    dict_states = build_dict_states(t_grid, cfg)
    D0_clean, phys_clean, K_clean = build_core(t_grid, dict_states, psi0, cfg)
    exact_states = exact_time_evolution(t_grid, psi0, cfg)

    # 先做整体的“噪声→误差”扫描
    noise_list = cfg["noise_sigma_list"]
    N_TRIALS = cfg["num_trials_per_sigma"]
    base_seed = cfg["base_seed"]

    mean12, std12, mean11, std11 = [], [], [], []

    for sigma in noise_list:
        errs12, errs11 = [], []
        for trial in range(N_TRIALS):
            np.random.seed(base_seed + 1000 * int(sigma*1000) + trial)
            e12, e11 = run_once(sigma, dict_states, D0_clean, phys_clean, K_clean,
                                psi0, exact_states, cfg, record_history=False)
            errs12.append(e12)
            errs11.append(e11)
        m12, s12 = np.mean(errs12), np.std(errs12)
        m11, s11 = np.mean(errs11), np.std(errs11)
        mean12.append(m12); std12.append(s12)
        mean11.append(m11); std11.append(s11)
        print(f"[{cfg['noise_injection_mode']}] sigma={sigma:.3f} | Eq.(12): {m12:.4f}±{s12:.4f} | Eq.(11): {m11:.4f}±{s11:.4f}")

    # 图1：噪声→误差
    '''plt.figure(figsize=(6,4))
    plt.errorbar(noise_list, mean12, yerr=std12, fmt="o-", capsize=4, label="Eq.(12)")
    plt.errorbar(noise_list, mean11, yerr=std11, fmt="s-", capsize=4, label="Eq.(11)")
    plt.xlabel("noise sigma")
    plt.ylabel("mean state error")
    plt.title(f"Noise robustness ({cfg['noise_injection_mode']})")
    plt.legend()
    plt.tight_layout()
    # 保存图1（新增）
    #plt.savefig(SAVE_DIR / "noise_robustness.png", bbox_inches="tight", facecolor="white")
    #plt.savefig(SAVE_DIR / "noise_robustness.pdf", bbox_inches="tight", facecolor="white")
    #plt.show()'''

    # 再单独挑一档噪声 -> 画 loss 和 observable
    sigma_show = cfg["plot_sigma"]
    np.random.seed(base_seed + 777)
    e12, e11, approx12, approx11, hist12, hist11 = run_once(
        sigma_show, dict_states, D0_clean, phys_clean, K_clean,
        psi0, exact_states, cfg, record_history=True
    )

    # 图2：loss 收敛
    '''plt.figure(figsize=(6,4))
    plt.plot(hist12, label=f"Eq.(12) loss, σ={sigma_show}")
    plt.plot(hist11, label=f"Eq.(11) loss, σ={sigma_show}")
    plt.xlabel("iteration")
    plt.ylabel("loss")
    plt.yscale("log")
    plt.legend()
    plt.title("Loss vs iteration")
    plt.tight_layout()'''
    # 保存图2（新增）
    #plt.savefig(SAVE_DIR / "loss_vs_iter4.png", bbox_inches="tight", facecolor="white")
    #plt.savefig(SAVE_DIR / "loss_vs_iter.pdf", bbox_inches="tight", facecolor="white")
    #plt.show()

    # 图3：可观测量（ΣZ）对比
    Z_sum = build_observable_Zsum(cfg)
    exact_Z = exp_O(exact_states, Z_sum)
    approxZ12 = exp_O(approx12, Z_sum)
    approxZ11 = exp_O(approx11, Z_sum)

    plt.figure(figsize=(6,4))
    plt.plot(t_grid, exact_Z.real, label="Exact ⟨ΣZ⟩")
    plt.plot(t_grid, approxZ12.real, "--", label=f"Eq.(12) ⟨ΣZ⟩, σ={sigma_show}")
    plt.plot(t_grid, approxZ11.real, ":", label=f"Eq.(11) ⟨ΣZ⟩, σ={sigma_show}")
    plt.xlabel("time")
    plt.ylabel("⟨ΣZ⟩")
    plt.title("Observable comparison")
    plt.legend()
    plt.tight_layout()
    # 保存图3（新增）
    plt.savefig(SAVE_DIR / "observable_comparison4.png", bbox_inches="tight", facecolor="white")
    #plt.savefig(SAVE_DIR / "observable_comparison.pdf", bbox_inches="tight", facecolor="white")
    plt.show()

    print(f"[info] system qubits = {n}, NISQ大约需要 {n+1} qubits")